/**
 * @file DiskInfo.c
 * 
 * @author https://ay123.net 
 * @version 0.1
 * @date 2024-09-08
 * 
 * @copyright Copyright (c) 2015 - 2024
 * 
 */

#include <Uefi.h>
#include <Library/PcdLib.h>
#include <Library/UefiLib.h>
#include <Library/UefiApplicationEntryPoint.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/MemoryAllocationLib.h>

#include <Protocol/BlockIo.h>
#include <Protocol/DiskIo.h>
#include <Protocol/DiskInfo.h>
#include <Library/BaseMemoryLib.h>
#include <Library/PrintLib.h>
#include <Protocol/IdeControllerInit.h>

VOID HexDump (UINT8 *Buffer, UINT8 RowNum)
{
  UINT8   Cols;
  UINT8   Rows;

  for (Rows = 0; Rows < RowNum; Rows ++) {
    for (Cols = 0; Cols < 16; Cols ++) {
      Print (L"%2X ", Buffer[Cols+Rows*16]);
    }

    Print (L"  ");

    for (Cols = 0; Cols < 16; Cols ++) {
      if ((Buffer[Cols+Rows*16] >= '0' && Buffer[Cols+Rows*16] <= '9') ||
         (Buffer[Cols+Rows*16] >= 'a' && Buffer[Cols+Rows*16] <= 'z') ||
         (Buffer[Cols+Rows*16] >= 'A' && Buffer[Cols+Rows*16] <= 'Z')) 
        Print (L"%c", Buffer[Cols+Rows*16]);
      else 
        Print (L".");
    }
    Print (L"\n\r");
  }
  Print (L"\n\r");
}

/**
  Eliminate the extra spaces in the Str to one space.

  @param    Str     Input string info.
**/
VOID
BmEliminateExtraSpaces (
  IN CHAR16                    *Str
  )
{
  UINTN                        Index;
  UINTN                        ActualIndex;

  for (Index = 0, ActualIndex = 0; Str[Index] != L'\0'; Index++) {
    if ((Str[Index] != L' ') || ((ActualIndex > 0) && (Str[ActualIndex - 1] != L' '))) {
      Str[ActualIndex++] = Str[Index];
    }
  }
  Str[ActualIndex] = L'\0';
}

EFI_STATUS
EFIAPI
GetDiskIdentifyData (EFI_HANDLE Handle)
{
  EFI_STATUS                  Status;
  EFI_DISK_INFO_PROTOCOL      *DiskInfo;
  EFI_ATAPI_IDENTIFY_DATA     IdentifyData;
  UINT32                      IdentifyDataSize;
  CHAR16                      *ModelName;
  CHAR16                      *SerialNo;
  CONST UINTN                 ModelNameLength = 40;
  CONST UINTN                 SerialNoLength  = 20;
  UINTN                       Index;

  Status = gBS->HandleProtocol(Handle, &gEfiDiskInfoProtocolGuid, (VOID**)&DiskInfo);
  if (EFI_ERROR(Status)) {
      return Status;
  }

  //            
  // AHCI or IDE
  //
  if (CompareGuid (&DiskInfo->Interface, &gEfiDiskInfoAhciInterfaceGuid) ||
      CompareGuid (&DiskInfo->Interface, &gEfiDiskInfoIdeInterfaceGuid)) {

    IdentifyDataSize = sizeof (EFI_ATAPI_IDENTIFY_DATA);
    Status = DiskInfo->Identify(DiskInfo, &IdentifyData, &IdentifyDataSize);
    if (!EFI_ERROR (Status)) {

      ModelName = AllocatePool(sizeof(CHAR16) * ModelNameLength);
      SerialNo  = AllocatePool(sizeof(CHAR16) * SerialNoLength);

      // According to the ATA specification, the model name and serial number fields 
      // in the identify data are stored as an array of 16-bit words.
      // Each word is stored in little-endian format, meaning the least significant byte comes first.
      
      for (Index = 0; Index + 1 < ModelNameLength; Index += 2) {
        ModelName[Index]     = (CHAR16) IdentifyData.ModelName[Index + 1];
        ModelName[Index + 1] = (CHAR16) IdentifyData.ModelName[Index];
      }

      for (Index = 0; Index + 1 < SerialNoLength; Index += 2) {
        SerialNo[Index]     = (CHAR16) IdentifyData.SerialNo[Index + 1];
        SerialNo[Index + 1] = (CHAR16) IdentifyData.SerialNo[Index];
      }

      BmEliminateExtraSpaces(ModelName);
      BmEliminateExtraSpaces(SerialNo);

      HexDump((UINT8 *)&IdentifyData, 16);

      Print (L"Model Name: %s\n\r", ModelName);
      Print (L"Serial No : %s\n\r", SerialNo);
      Print (L"Disk type: %g\n\r", DiskInfo->Interface);
    }
    else {
      return Status;
    }
  }

  return EFI_SUCCESS;
}

/**
  The user Entry Point for Application. The user code starts with this function
  as the real entry point for the application.

  @param[in] ImageHandle    The firmware allocated handle for the EFI image.
  @param[in] SystemTable    A pointer to the EFI System Table.

  @retval EFI_SUCCESS       The entry point is executed successfully.
  @retval other             Some error occurs when executing this entry point.

**/
EFI_STATUS
EFIAPI
UefiMain (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS                   Status;
  UINTN                        HandleCount = 0;
  EFI_HANDLE                   *HandleBuffer = NULL;
  UINTN                        Index;
  EFI_BLOCK_IO_PROTOCOL       *BlockIo;
  EFI_DISK_IO_PROTOCOL        *DiskIo;

  Status = gBS->LocateHandleBuffer(
    ByProtocol,
    &gEfiBlockIoProtocolGuid,
    NULL,
    &HandleCount,
    &HandleBuffer
  );

  if (EFI_ERROR(Status)) {
    Print(L"Failed to locate block I/O handles: %r\n", Status);
    return Status;
  }

  for (Index = 0; Index < HandleCount; Index++) {
    Status = gBS->HandleProtocol(
      HandleBuffer[Index],
      &gEfiBlockIoProtocolGuid,
      (VOID**)&BlockIo
    );

    if (EFI_ERROR(Status) || BlockIo == NULL || BlockIo->Media == NULL) {
      continue;
    }

    Status = gBS->HandleProtocol(
      HandleBuffer[Index],
      &gEfiDiskIoProtocolGuid,
      (VOID**)&DiskIo
    );

    if (!EFI_ERROR(Status) && DiskIo != NULL && !BlockIo->Media->RemovableMedia) {

      if (BlockIo->Media->LogicalPartition == FALSE &&
          BlockIo->Media->BlockSize > 0 &&
          BlockIo->Media->LastBlock > 0) {
          GetDiskIdentifyData(HandleBuffer[Index]);
          Print (L"Size : %d MB\n\r", (BlockIo->Media->LastBlock + 1) * BlockIo->Media->BlockSize / (1024 * 1024));
          Print(L"\n\r");
      }
    }
  }

  if (HandleBuffer != NULL) {
    gBS->FreePool(HandleBuffer);
  }

  return EFI_SUCCESS;
}